import torch
import numpy as np
import random
import tqdm


def get_rolling_window_multistep(forecasting_length, interval_length, window_length, features, labels):
    """
    修改后的滑动窗口函数 - 使预测从当前天开始
    参数:
        forecasting_length: 预测的未来时间步数（输出窗口长度）
        interval_length: 输入窗口与预测窗口之间的间隔步数（设置为0以实现连续预测）
        window_length: 输入窗口的历史时间步数（输入窗口长度）
        features: 输入特征数据，形状为 (特征数, 时间步数)
        labels: 标签数据，形状需与 features 的时间维度一致
    返回:
        输入序列 (n_samples, n_features, window_length)
        输出序列 (n_samples, 1, forecasting_length)
    """
    # 确保输入数据形状正确
    if features.shape[1] != labels.shape[1]:
        raise ValueError("Features and labels must have the same time dimension")

    # 计算可生成的样本数量
    n_samples = features.shape[1] - window_length - forecasting_length - interval_length + 1

    # 初始化输出数组
    output_features = np.zeros((n_samples, features.shape[0], window_length))
    output_labels = np.zeros((n_samples, 1, forecasting_length))

    # 生成滑动窗口
    for index in tqdm.tqdm(range(n_samples), desc='Preparing data windows'):
        # 输入窗口：[index, index+window_length)
        output_features[index] = features[:, index:index + window_length]

        # 关键修改点：输出窗口从输入窗口的最后一天开始 [index+window_length-1, index+window_length-1+forecasting_length)
        output_labels[index] = labels[:,
                               index + window_length:
                               index + window_length + forecasting_length]

    return torch.from_numpy(output_features), torch.from_numpy(output_labels)


def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.benchmark = True
    torch.backends.cudnn.deterministic = False